package org.ala.tiktools.sharddb.mybatis.executor;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.mybatis.executor.SqlSessionTemplateExecutor;
import org.ala.tiktools.mybatis.pager.Pager;
import org.ala.tiktools.sharddb.db.ThreadLocalLookupKey;
import org.ala.tiktools.sharddb.exception.ShardDBRuntimeException;
import org.ala.tiktools.sharddb.mybatis.interceptor.Shard;
import org.ala.tiktools.sharddb.mybatis.interceptor.ShardInfoThreadLocal;
import org.ala.tiktools.sharddb.mybatis.router.ShardInfo;
import org.ala.tiktools.tools.kv.KV;
import org.ala.tiktools.tools.kv.Maps;
import org.slf4j.Logger;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @author ala
 * @date 2025-03-09 15:39
 */
public class ShardSqlSessionTemplateExecutor {


    static Logger log = LoggerFactory.getLogger("shard", "SqlSession");


    static ShardSqlSessionTemplateExecutor instance;
    /**
     *  单例
     */
    public static ShardSqlSessionTemplateExecutor instance() {
        if (instance == null) {
            synchronized (ShardSqlSessionTemplateExecutor.class) {
                if (instance == null) {
                    instance = new ShardSqlSessionTemplateExecutor();
                }
            }
        }
        return instance;
    }


    /**
     *  分库分表模式执行方法
     */
    public <E> Map<ShardInfo, E> selectOne(Shard anno, ThreadPoolExecutor threadPool, Method method, Map<ShardInfo, Object[]> args) {
        CountDownLatch latch = new CountDownLatch(args.size());
        Map<ShardInfo, E> res = new ConcurrentHashMap<>();
        Map<ShardInfo, Throwable> errors = new HashMap<>();

        long now = System.currentTimeMillis();
        args.forEach(((shardInfo, objects) -> {
            try {
                threadPool.execute(() -> {
                    try {
                        ThreadLocalLookupKey.setLookup(shardInfo.getShardDB());
//                        ShardInfoThreadLocal.set(shardInfo);

                        ShardExecutor. Result r = ShardExecutor.doInvoke(anno, shardInfo, method, objects);
//                        E e = SqlSessionTemplateExecutor.instance().selectOne(method, objects, Maps.build(KV.build("shardSuffix", shardInfo.getShardTable())));
                        if (!r.isSuccess()) {errors.put(shardInfo, r.getE()); return;}
                        res.put(shardInfo, (E) r.getRes());
                    } catch (Exception e) {
                        errors.put(shardInfo, e);
                    } finally {
                        ThreadLocalLookupKey.clear();
//                        ShardInfoThreadLocal.clear();
                        latch.countDown();
                    }
                });
            } catch (Exception e) {
                latch.countDown();
                throw new ShardDBRuntimeException("[ShardSqlSessionTemplateExecutor selectOne] 线程池满，无法执行sql任务", e);
            }
        }));
        try {
            latch.await(anno.timeout(), TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new ShardDBRuntimeException(e);
        }
        //  如果有任何人执行失败，直接抛异常
        if (!errors.isEmpty()) {
            throw new ShardDBRuntimeException(errors.values().iterator().next());
        }
        //  检测如果有人没有执行完成，抛超时异常，或者有人执行失败，抛失败异常
        args.forEach(((shardInfo, objects) -> {
            if (!res.containsKey(shardInfo)) {
                throw new ShardDBRuntimeException(String.format("[ShardSqlSessionTemplateExecutor selectOne] sql任务执行超时 db:%s table:%s method:%s class:%s",
                            shardInfo.getShardDB(), shardInfo.getShardTable(),
                            method.getName(), method.getDeclaringClass().getSimpleName()
                        ));
            }
            if (errors.containsKey(shardInfo)) {
                throw new ShardDBRuntimeException(errors.get(shardInfo));
            }
        }));

//        log.info(String.format("[ShardSqlSessionTemplateExecutor selectOne] method:%s shardSize:%d delay:%d", method.getName(), args.size(), System.currentTimeMillis() - now));

        return res;
    }


    /**
     *  分库分表模式执行方法
     */
    public <E> Map<ShardInfo, List<E>> selectList(Shard anno, ThreadPoolExecutor threadPool, Method method, Map<ShardInfo, Object[]> args) {
        CountDownLatch latch = new CountDownLatch(args.size());
        Map<ShardInfo, List<E>> res = new ConcurrentHashMap<>();
        Map<ShardInfo, Throwable> errors = new HashMap<>();

        long now = System.currentTimeMillis();
        args.forEach(((shardInfo, objects) -> {
            try {
                threadPool.execute(() -> {
                    try {
                        ThreadLocalLookupKey.setLookup(shardInfo.getShardDB());
//                        ShardInfoThreadLocal.set(shardInfo);

                        ShardExecutor. Result r = ShardExecutor.doInvoke(anno, shardInfo, method, objects);
//                        E e = SqlSessionTemplateExecutor.instance().selectOne(method, objects, Maps.build(KV.build("shardSuffix", shardInfo.getShardTable())));
                        if (!r.isSuccess()) {errors.put(shardInfo, r.getE()); return;}
                        List datas = null;
                        if (r.getRes() instanceof Pager) {
                            datas = ((Pager)r.getRes()).getDatas();
                        } else if (r.getRes() instanceof List) {
                            datas = (List) r.getRes();
                        }
                        if (datas == null) {datas = new ArrayList();}
                        res.put(shardInfo, datas);
                    } catch (Exception e) {
                        errors.put(shardInfo, e);
                    } finally {
                        ThreadLocalLookupKey.clear();
//                        ShardInfoThreadLocal.clear();
                        latch.countDown();
                    }
                });
            } catch (Exception e) {
                latch.countDown();
                throw new ShardDBRuntimeException("[ShardSqlSessionTemplateExecutor selectList] 线程池满，无法执行sql任务", e);
            }
        }));
        try {
            long timeout = anno.timeout();
            latch.await(timeout, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new ShardDBRuntimeException(e);
        }
        if (!errors.isEmpty()) {
            throw new ShardDBRuntimeException(errors.values().iterator().next());
        }
        //  检测如果有人没有执行完成，抛超时异常，或者有人执行失败，抛失败异常
        args.forEach(((shardInfo, objects) -> {
            if (!res.containsKey(shardInfo)) {
                throw new ShardDBRuntimeException(String.format("[ShardSqlSessionTemplateExecutor selectList] sql任务执行超时 db:%s table:%s method:%s class:%s",
                        shardInfo.getShardDB(), shardInfo.getShardTable(),
                        method.getName(), method.getDeclaringClass().getSimpleName()
                ));
            }
            if (errors.containsKey(shardInfo)) {
                throw new ShardDBRuntimeException(errors.get(shardInfo));
            }
        }));

//        log.info(String.format("[ShardSqlSessionTemplateExecutor selectOne] method:%s shardSize:%d delay:%d", method.getName(), args.size(), System.currentTimeMillis() - now));


        return res;
    }



}
